#pragma once

#include "HALS.h"  // for ALS()

namespace factor {
  template <typename MAT> class StochHALS : public HALS<MAT> {
  public:
    struct Iteration {
      int iter;
      double mse;
      posix::time_duration duration;

    Iteration(int i, double mse_, posix::time_duration d)
    : iter(i), mse(mse_), duration(d) {}

      void Print() {
	MSG("iter=" << iter << " rmse=" << mse
	    << " d=" << SECONDS(duration));
      }

      void Output(ostream& out) {
	out << iter << ",\t" << SECONDS(duration) << ",\t" << mse 
	    << endl;
      }    
    }; // class Iteration

    vector<Iteration> progress;
    vector<int> r, c; // start row/col of each block
    int row_block, col_block;
    double mse, oldMse, currentMse, sMse;
    double step;

    int k_rows, k_cols, s_row, s_col;
    int n_sweep;
    
    StochHALS(mpi::communicator& world_, po::variables_map& vm_) 
      : HALS<MAT>(world_, vm_) {}
    
    void NMF(MAT& A) {
      this->startTime = CLOCK;
      this->progress.clear();

      ARMA<MAT>::getSize(A, this->par.m, this->par.n);
      this->ProcessParameter();
    
      this->W = randu<mat>(this->par.m, this->par.k); // initialize
      this->H = randu<mat>(this->par.n, this->par.k);
    
      this->normA = norm(A, "fro");
      this->stop = false;
      n_sweep = 1;
      k_rows = this->par.m / n_sweep;
      k_cols = this->par.n / n_sweep;

      mse = 1e-6;
      sMse = 0;
      int count = 1;
      step = 1;
      FORE(iter, 0, this->par.maxiter) {
	PickRowsCols();
	OptimizeRowsCols(A);
	sMse += currentMse;
	
	if (iter % n_sweep == n_sweep-1) {
	  oldMse = mse;
	  mse = sMse / n_sweep;
	  sMse = 0;
  
	  this->stop = mse < this->par.tol || fabs(mse-oldMse) < this->par.tol
	      || fabs(mse-oldMse)/oldMse < this->par.tol;
	  
	  Iteration info(iter, mse, CLOCK-this->startTime);
	  progress.push_back(info);
	  if (this->PrintCondition(iter)) info.Print();
	  if (!this->par.output.empty()) {
	    ofstream f(this->par.output.c_str(), ios::app);
	    info.Output(f);
	  }

	  if (mse > oldMse) {
	    step = pow( 1.0 / (++count), this->par.alpha);
	    if (step < 1e-6) count = 1;
	    MSG("step=" << step);
	  }	  
	}

	if (this->stop) break;
      } // for iter

      if (!this->par.output.empty()) {     // output factors
	this->W.save(this->par.output+"_W");
	this->H.save(this->par.output+"_H");
      }
      
    }

    void PickRowsCols() {
      s_row = rand() % (this->par.m - k_rows + 1);
      s_col = rand() % (this->par.n - k_cols + 1);
      //MSG(s_row << " " << s_row+k_rows-1 << " " << s_col << " " << s_col+k_cols-1);
    }

    void OptimizeRowsCols(MAT &A) {
      // s_row .. s_row+k_rows-1
      // s_col .. s_col+k_cols-1
#define Asub A.submat(s_row, s_col, s_row+k_rows-1, s_col+k_cols-1)
#define Hsub this->H.rows(s_col, s_col+k_cols-1)
#define Wsub this->W.rows(s_row, s_row+k_rows-1)

      mat AH = Asub*Hsub, HTH = trans(Hsub)*Hsub;   // HALS for each column of W
      FORE(i, 0, this->par.k) {
	Wsub.col(i) = ALS(AH, HTH, Wsub, i, step);
      }
      
      AH = trans(Asub)*Wsub, HTH = trans(Wsub)*Wsub; // HALS for each column of H
      FORE(i, 0, this->par.k) {
	Hsub.col(i) = ALS(AH, HTH, Hsub, i, step);
      }
      
      double f = norm(Asub-Wsub*trans(Hsub),"fro"); f *= f;
      f = f / k_rows / k_cols;
      currentMse = sqrt(f);
    }

    void NMF(MAT& Arows, MAT& Acols) {
      StartNMF(Arows, Acols);
      
      sMse = 0;
      mse = 1e-6;
      int count = 1;
      step = 1;
      FORE(iter, 0, this->par.maxiter) {
	//step = pow( 1.0 / (count++), this->par.alpha);
	//if (step < 1e-6) count = 1;
	OptimizeBlock(Arows, Acols);
	SendWHBlock();

	if (MPI_IS_MASTER) {
	  sMse += currentMse;
	  if (MPI_SIZE == 1 || iter % MPI_SIZE == MPI_SIZE-1) {	    
	    oldMse = mse;
	    mse = sMse / MPI_SIZE;
	    sMse = 0;
	    
	    Iteration info(iter, mse, CLOCK-this->startTime);
	    progress.push_back(info);
	    if (this->PrintCondition(iter)) info.Print();
	    if (!this->par.output.empty()) {
	      ofstream f(this->par.output.c_str(), ios::app);
	      info.Output(f);
	    }

	    this->stop = mse < this->par.tol || fabs(mse-oldMse) < this->par.tol
	      || fabs(mse-oldMse)/oldMse < this->par.tol;

	    if (mse > oldMse) {
	      step = pow( 1.0 / (++count), this->par.alpha);
	      //if (step < 1e-6) count = 1;
	    }
	  }
	}
	broadcast(MPI_WORLD, this->stop, 0);
	broadcast(MPI_WORLD, step, 0);
	if (this->stop) break;
      }
    }

    void OptimizeBlock(MAT& Arows, MAT& Acols) {
      int fr = r[row_block], lr = r[row_block+1]-1;
      mat AH, HTH;

      HTH = trans(this->H)*this->H;   // HALS for each column of W
      if (MPI_SIZE > 1) AH = Acols.rows(fr,lr)*this->H;
      else AH = Acols*this->H; 
      FORE(i, 0, this->par.k) {
	this->W.col(i) = ALS(AH, HTH, this->W, i, step);
      }
      
      HTH = trans(this->W)*this->W;
      if (MPI_SIZE > 1) AH = trans(Acols.rows(fr,lr))*this->W; // HALS for each column of H
      else AH = trans(Acols)*this->W;
      FORE(i, 0, this->par.k) {
	this->H.col(i) = ALS(AH, HTH, this->H, i, step);
      }
      
      double f = norm(Acols.rows(fr,lr) - this->W*trans(this->H), "fro"), e = f*f, se;
      double n = (double)(lr-fr+1) * Acols.n_cols, sn;
      reduce(MPI_WORLD, e, se, std::plus<double>(), 0);
      reduce(MPI_WORLD, n, sn, std::plus<double>(), 0);
      if (MPI_IS_MASTER) currentMse = sqrt(se / sn);
    }

    void CheckStopCondition() {
      if (MPI_IS_MASTER) {
	this->stop = mse < this->par.tol || fabs(mse-oldMse) < this->par.tol
	  || fabs(mse-oldMse)/oldMse < this->par.tol;
      }
      broadcast(MPI_WORLD, this->stop, 0);
    }

#define STOCH_HALS_ROW_TAG 2
#define STOCH_HALS_MAT_TAG 3

    void SendWHBlock() {
      if (MPI_SIZE == 1) return;
      int p_send = MPI_RANK+1 >= MPI_SIZE ? 0 : MPI_RANK+1;
      int p_recv = MPI_RANK-1 < 0 ? MPI_SIZE-1 : MPI_RANK-1;

      vector<int> sidx(MPI_SIZE), ridx(MPI_SIZE);
      if (MPI_IS_MASTER) {
	RandPerm(sidx);
	FORE(i, 0, MPI_SIZE) ridx[sidx[i]] = i;
      }
      broadcast(MPI_WORLD, sidx, 0);
      broadcast(MPI_WORLD, ridx, 0);

      p_send = sidx[MPI_RANK];
      p_recv = ridx[MPI_RANK];
      int new_row;
      mat new_W;
      vector<mpi::request> reqs;
      reqs.push_back(this->world.isend(p_send, STOCH_HALS_ROW_TAG, row_block));
      reqs.push_back(this->world.isend(p_send, STOCH_HALS_MAT_TAG, this->W));
      reqs.push_back(this->world.irecv(p_recv, STOCH_HALS_ROW_TAG, new_row));
      reqs.push_back(this->world.irecv(p_recv, STOCH_HALS_MAT_TAG, new_W));
      mpi::wait_all(ALL(reqs));
      row_block = new_row;
      this->W = new_W;
    }

    void RandPerm(vector<int>& sidx) {
      FORE(i, 0, MPI_SIZE) sidx[i] = i;
      FORE(i, 0, MPI_SIZE) {
	int pos = rand() % (MPI_SIZE-i) + i;
	swap(sidx[i], sidx[pos]);
      }
    }

    void StartNMF(MAT& Arows, MAT& Acols) {
      this->startTime = CLOCK;
      this->progress.clear();

      ARMA<MAT>::getSize(Arows, this->mr, this->par.n);
      ARMA<MAT>::getSize(Acols, this->par.m, this->nc);
      vector<int> rs, cs;
      gather(this->world, this->mr, rs, 0);
      gather(this->world, this->nc, cs, 0);
      if (MPI_IS_MASTER) {
	r = c = vector<int>(SIZE(this->world)+1, 0);
	FORI(i, 1, SIZE(this->world)) {
	  r[i] = r[i-1]+rs[i-1];
	  c[i] = c[i-1]+cs[i-1];
	}
	if (r.back() != this->par.m || c.back() != this->par.n) {
	  FATAL("Check size of blocks");
	}
      }
      broadcast(this->world, r, 0);
      broadcast(this->world, c, 0);
      
      this->ProcessParameter();    
      this->ParInitialize();
      this->stop = false;

      row_block = this->world.rank();
      col_block = this->world.rank();
    }

    void BlockSize(int i, int j, int &rs, int& cs) {
      rs = r[i+1]-r[i];
      cs = c[i+1]-c[i];
    }
  };
}
